from torch.optim import AdamW

from demo5 import MyModel
from demo4 import MyDataset
import torch
from torch.utils.data import DataLoader
from transformers import BertTokenizer

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCH = 2
token = BertTokenizer.from_pretrained(r"D:\pythonWork\PythonProject3\model\google-bert\bert-base-chinese\models--google-bert--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")
train_dataset = MyDataset("test")
#自定义数据集的collate_fn 对数据进行编码处理
def collate_fn(data):
    sentes = [ i[0] for i in data ]
    label = [ i[1] for i in data ]
    data = token.batch_encode_plus(
        batch_text_or_text_pairs=sentes,
        truncation=True,
        padding="max_length",
        max_length=350,
        return_tensors="pt",
        return_length = True
    )
    input_ids =  data["input_ids"]
    attention_mask = data["attention_mask"]
    token_type_ids = data["token_type_ids"]
    label = torch.LongTensor(label)
    return input_ids, attention_mask, token_type_ids,label

train_loader = DataLoader(train_dataset,
                        batch_size=32,
                        shuffle=True,
                        drop_last=True,
                        collate_fn=collate_fn
                        )
if __name__ == '__main__':
    acc = 0
    total = 0
    print(DEVICE)
    model = MyModel().to(DEVICE)
    model.load_state_dict(torch.load(r"params/0bert.pt"))
    # 模型训练模式
    model.eval()
    for epoch in range(EPOCH):
        for i, (input_ids, attention_mask, token_type_ids, label) in enumerate(train_loader):
            #将数据放到设备上
            input_ids = input_ids.to(DEVICE)
            attention_mask = attention_mask.to(DEVICE)
            token_type_ids = token_type_ids.to(DEVICE)
            label = label.to(DEVICE)
            #执行前向计算得到输出
            output = model(input_ids, attention_mask, token_type_ids)
            out = output.argmax(dim=1)
            acc += (out == label).sum().item()
            total+= len(label)
            print("acc:",acc/total)
